#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <io.h>
#include <fcntl.h>
#include <direct.h>
#include <windows.h>

// 设置文件模式为二进制
void set_binary_mode(FILE* file) {
#ifdef _WIN32
    _setmode(_fileno(file), _O_BINARY);
#endif
}

// 确保路径中的反斜杠被正确处理
void sanitize_path(char* path) {
    for (char* p = path; *p; p++) {
        if (*p == '/') {
            *p = '\\';  // 将正斜杠转换为反斜杠
        }
    }
}

// RLE压缩函数
unsigned char* rle_compress(const unsigned char* data, size_t size, size_t* compressed_size) {
    if (size == 0) {
        *compressed_size = 0;
        return NULL;
    }

    unsigned char* compressed = (unsigned char*)malloc(size * 2);
    if (!compressed) {
        *compressed_size = 0;
        return NULL;
    }

    unsigned char* dest = compressed;
    const unsigned char* src = data;
    const unsigned char* end = data + size;

    while (src < end) {
        unsigned char byte = *src++;
        int count = 1;

        while (src < end && *src == byte && count < 255) {
            src++;
            count++;
        }

        *dest++ = (unsigned char)count;
        *dest++ = byte;
    }

    *compressed_size = dest - compressed;
    return compressed;
}

// 压缩单个文件
int compress_file(const char* input_path, const char* output_path, const char* rel_path) {
    FILE* input_file = fopen(input_path, "rb");
    if (!input_file) {
        perror("无法打开输入文件");
        return 1;
    }
    set_binary_mode(input_file);

    fseek(input_file, 0, SEEK_END);
    size_t file_size = ftell(input_file);
    fseek(input_file, 0, SEEK_SET);

    unsigned char* file_data = (unsigned char*)malloc(file_size);
    if (!file_data) {
        perror("内存分配失败");
        fclose(input_file);
        return 1;
    }
    fread(file_data, 1, file_size, input_file);
    fclose(input_file);

    size_t compressed_size;
    unsigned char* compressed_data = rle_compress(file_data, file_size, &compressed_size);
    free(file_data);

    if (!compressed_data) {
        perror("压缩失败");
        return 1;
    }

    FILE* output_file = fopen(output_path, "ab");
    if (!output_file) {
        perror("无法打开输出文件");
        free(compressed_data);
        return 1;
    }
    set_binary_mode(output_file);

    fputc(1, output_file); // 文件类型标记

    size_t filename_len = strlen(rel_path);
    fwrite(&filename_len, sizeof(size_t), 1, output_file);
    fwrite(rel_path, 1, filename_len, output_file);

    fwrite(&file_size, sizeof(size_t), 1, output_file);
    fwrite(&compressed_size, sizeof(size_t), 1, output_file);

    fwrite(compressed_data, 1, compressed_size, output_file);

    fclose(output_file);
    free(compressed_data);
    return 0;
}

// 判断路径是否为目录
int is_directory(const char* path) {
    DWORD file_attr = GetFileAttributesA(path);
    return (file_attr != INVALID_FILE_ATTRIBUTES &&
        (file_attr & FILE_ATTRIBUTE_DIRECTORY));
}

// 压缩目录及其子内容
void compress_directory(const char* dir_path, const char* output_path, const char* base_path) {
    char rel_path[MAX_PATH] = "";

    // 计算相对于base_path的相对路径
    if (base_path && strlen(base_path) > 0) {
        if (strcmp(dir_path, base_path) == 0) {
            // 根目录：使用basename作为相对路径（如"1"）
            const char* last_slash = strrchr(base_path, '\\');
            if (last_slash) {
                strcpy(rel_path, last_slash + 1);
            }
            else {
                strcpy(rel_path, base_path);
            }
            // 确保根目录路径以斜杠结尾（如"1\"）
            if (strlen(rel_path) > 0 && rel_path[strlen(rel_path) - 1] != '\\') {
                strncat(rel_path, "\\", MAX_PATH - strlen(rel_path) - 1);
            }
        }
        else {
            // 子目录：计算完整相对路径（如"1\2\3\"）
            size_t base_len = strlen(base_path);
            if (strncmp(dir_path, base_path, base_len) == 0) {
                // 提取相对部分（如"\2\3"）
                const char* rel_part = dir_path + base_len;
                // 复制相对部分并移除前导斜杠
                if (rel_part[0] == '\\') rel_part++;

                // 关键修复：使用完整的相对路径
                const char* base_name = strrchr(base_path, '\\');
                if (base_name) {
                    base_name++;  // 指向basename的第一个字符
                    snprintf(rel_path, MAX_PATH, "%s\\%s", base_name, rel_part);
                }
                else {
                    strncpy(rel_path, rel_part, MAX_PATH - 1);
                }

                // 确保目录路径以斜杠结尾
                if (strlen(rel_path) > 0 && rel_path[strlen(rel_path) - 1] != '\\') {
                    strncat(rel_path, "\\", MAX_PATH - strlen(rel_path) - 1);
                }
            }
        }
    }

    // 确保路径格式正确
    sanitize_path(rel_path);

    // 写入目录条目头部
    printf("压缩目录: %s (相对路径: %s)\n", dir_path, rel_path);

    FILE* output_file = fopen(output_path, "ab");
    if (!output_file) {
        perror("无法打开输出文件");
        return;
    }
    set_binary_mode(output_file);

    fputc(2, output_file);  // 目录类型标记

    size_t path_len = strlen(rel_path);
    fwrite(&path_len, sizeof(size_t), 1, output_file);
    fwrite(rel_path, 1, path_len, output_file);

    fclose(output_file);

    // 递归处理子项
    char search_path[MAX_PATH];
    sprintf(search_path, "%s\\*", dir_path);

    WIN32_FIND_DATAA find_data;
    HANDLE hFind = FindFirstFileA(search_path, &find_data);

    if (hFind == INVALID_HANDLE_VALUE) {
        printf("无法打开目录: %s\n", dir_path);
        return;
    }

    do {
        if (strcmp(find_data.cFileName, ".") == 0 ||
            strcmp(find_data.cFileName, "..") == 0) {
            continue;
        }

        char child_path[MAX_PATH];
        sprintf(child_path, "%s\\%s", dir_path, find_data.cFileName);

        if (find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
            compress_directory(child_path, output_path, base_path);
        }
        else {
            // 构建文件的相对路径（如"1\2\3\file.txt"）
            char rel_file_path[MAX_PATH] = { 0 };
            if (strlen(rel_path) > 0) {
                snprintf(rel_file_path, MAX_PATH, "%s%s", rel_path, find_data.cFileName);
            }
            else {
                strcpy(rel_file_path, find_data.cFileName);
            }

            // 确保路径格式正确
            sanitize_path(rel_file_path);

            printf("压缩文件: %s (相对路径: %s)\n", child_path, rel_file_path);
            compress_file(child_path, output_path, rel_file_path);
        }
    } while (FindNextFileA(hFind, &find_data));

    FindClose(hFind);
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        printf("用法: rlezip [文件/目录1] [文件/目录2]... [输出文件.rzp]\n");
        return 1;
    }

    const char* output_path = argv[argc - 1];

    // 清空输出文件
    FILE* output_file = fopen(output_path, "wb");
    if (!output_file) {
        perror("无法创建输出文件");
        return 1;
    }
    fclose(output_file);

    for (int i = 1; i < argc - 1; i++) {
        const char* input_path = argv[i];

        if (is_directory(input_path)) {
            compress_directory(input_path, output_path, input_path);
        }
        else {
            // 处理单个文件，相对路径为文件名
            const char* filename = strrchr(input_path, '\\') ? strrchr(input_path, '\\') + 1 : input_path;
            compress_file(input_path, output_path, filename);
        }
    }

    printf("压缩完成\n");
    return 0;
}